{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "56341c66",
   "metadata": {},
   "source": [
    "# Chapter 2: Emitting Basic IR\n",
    "\n",
    "Now that we're familiar with our language and the AST, let's see how xDSL can\n",
    "help to compile Toy.\n",
    "\n",
    "## Introduction: Multi-Level Intermediate Representation\n",
    "\n",
    "xDSL leverages the MLIR representation of a program. MLIR specifies a text format of\n",
    "this representation, which is useful for debugging and interoperation. For example,\n",
    "all the text in this format you'll see in this tutorial can be executed with the \n",
    "Toy language as compiled in the (MLIR Toy Tutorial)[https://mlir.llvm.org/docs/Tutorials/Toy/].\n",
    "\n",
    "Let's take a quick look at the textual IR representation of our example program:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "29e15305",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "builtin.module {\n",
      "  \"toy.func\"() ({\n",
      "    %0 = \"toy.constant\"() {value = dense<[[1.000000e+00, 2.000000e+00, 3.000000e+00], [4.000000e+00, 5.000000e+00, 6.000000e+00]]> : tensor<2x3xf64>} : () -> tensor<2x3xf64>\n",
      "    %1 = \"toy.reshape\"(%0) : (tensor<2x3xf64>) -> tensor<2x3xf64>\n",
      "    %2 = \"toy.constant\"() {value = dense<[1.000000e+00, 2.000000e+00, 3.000000e+00, 4.000000e+00, 5.000000e+00, 6.000000e+00]> : tensor<6xf64>} : () -> tensor<6xf64>\n",
      "    %3 = \"toy.reshape\"(%2) : (tensor<6xf64>) -> tensor<3x2xf64>\n",
      "    %4 = \"toy.reshape\"(%3) : (tensor<3x2xf64>) -> tensor<2x3xf64>\n",
      "    %5 = \"toy.add\"(%1, %4) : (tensor<2x3xf64>, tensor<2x3xf64>) -> tensor<2x3xf64>\n",
      "    \"toy.print\"(%5) : (tensor<2x3xf64>) -> ()\n",
      "    \"toy.return\"() : () -> ()\n",
      "  }) {sym_name = \"main\", function_type = () -> ()} : () -> ()\n",
      "}"
     ]
    }
   ],
   "source": [
    "from __future__ import annotations\n",
    "\n",
    "from toy.compiler import parse_toy\n",
    "\n",
    "from xdsl.printer import Printer\n",
    "\n",
    "example = \"\"\"\n",
    "def main() {\n",
    "  var a<2, 3> = [[1, 2, 3], [4, 5, 6]];\n",
    "  var b<3, 2> = [1, 2, 3, 4, 5, 6];\n",
    "  var c<2, 3> = b;\n",
    "  var d = a + c;\n",
    "  print(d);\n",
    "}\n",
    "\"\"\"\n",
    "\n",
    "toy_0 = parse_toy(example)\n",
    "Printer().print_op(toy_0)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "54bb582b",
   "metadata": {},
   "source": [
    "As you might have noticed, some parts look very similar to the Toy program above, and some\n",
    "things are added in. Let's look at the structure of an operation in the MLIR output before\n",
    "taking a close look at exactly what's going on."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "6251171e",
   "metadata": {},
   "source": [
    "## IR Syntax In Detail\n",
    "\n",
    "MLIR is designed to be a completely extensible infrastructure; there is no\n",
    "closed set of attributes (think: constant metadata), operations, or types. MLIR\n",
    "supports this extensibility with the concept of Dialects. Dialects provide a grouping \n",
    "mechanism for abstraction under a unique `namespace`.\n",
    "\n",
    "In MLIR, `Operations` are the core unit of abstraction and computation, similar in many \n",
    "ways to LLVM instructions. Operations can have application-specific semantics and can be \n",
    "used to represent all of the core IR structures in LLVM: instructions, globals (like \n",
    "functions), modules, etc.\n",
    "\n",
    "Here is the MLIR assembly for the Toy `transpose` operations:\n",
    "\n",
    "```mlir\n",
    "%t_tensor = \"toy.transpose\"(%tensor): (tensor<2x3xf64>) -> tensor<3x2xf64>\n",
    "```\n",
    "\n",
    "Let's break down the anatomy of this MLIR operation:\n",
    "\n",
    "-   `%t_tensor`\n",
    "\n",
    "    *   The name given to the result defined by this operation (which includes\n",
    "        [a prefixed sigil to avoid collisions](https://mlir.llvm.org/docs/LangRef/#identifiers-and-keywords)).\n",
    "        An operation may define zero or more results (in the context of Toy, we\n",
    "        will limit ourselves to single-result operations), which are SSA values.\n",
    "        The name is used during parsing but is not persistent (e.g., it is not\n",
    "        tracked in the in-memory representation of the SSA value).\n",
    "\n",
    "-   `\"toy.transpose\"`\n",
    "\n",
    "    *   The name of the operation. It is expected to be a unique string, with\n",
    "        the namespace of the dialect prefixed before the \"`.`\". This can be read\n",
    "        as the `transpose` operation in the `toy` dialect.\n",
    "\n",
    "-   `(%tensor)`\n",
    "\n",
    "    *   A list of zero or more input operands (or arguments), which are SSA\n",
    "        values defined by other operations or referring to block arguments.\n",
    "\n",
    "\n",
    "-   `(tensor<2x3xf64>) -> tensor<3x2xf64>`\n",
    "\n",
    "    *   This refers to the type of the operation in a functional form, spelling\n",
    "        the types of the arguments in parentheses and the type of the return\n",
    "        values afterward.\n",
    "\n",
    "\n",
    "Shown here is the general form of an operation. As described above,\n",
    "the set of operations in MLIR is extensible. Operations are modeled\n",
    "using a small set of concepts, enabling operations to be reasoned\n",
    "about and manipulated generically. These concepts are:\n",
    "\n",
    "-   A name for the operation.\n",
    "-   A list of SSA operand values.\n",
    "-   A list of attributes.\n",
    "-   A list of types for result values.\n",
    "-   A source location for debugging purposes.\n",
    "-   A list of successors blocks (for branches, mostly).\n",
    "-   A list of regions (for structural operations like functions).\n",
    "\n",
    "## Defining Toy Operations\n",
    "\n",
    "Now that we have a `Toy` dialect, we can start defining the operations. This\n",
    "will allow for providing semantic information that the rest of the system can\n",
    "hook into. As an example, let's walk through the creation of a `toy.constant`\n",
    "operation. This operation will represent a constant value in the Toy language.\n",
    "\n",
    "```mlir\n",
    " %4 = \"toy.constant\"() {value = dense<[[1, 2, 3], [4, 5, 6]]> : tensor<2x3xf64>} : () -> tensor<2x3xf64>\n",
    "```\n",
    "\n",
    "This operation takes zero operands, a dense elements attribute named `value` \n",
    "to represent the constant value, and returns a single result of RankedTensorType. \n",
    "Let's take a look at the full definition and step through it in detail.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "96c2a0c0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "%0 = \"toy.constant\"() {value = dense<[[1.000000e+00, 2.000000e+00, 3.000000e+00], [4.000000e+00, 5.000000e+00, 6.000000e+00]]> : tensor<2x3xf64>} : () -> tensor<2x3xf64>"
     ]
    }
   ],
   "source": [
    "from typing import TypeAlias\n",
    "\n",
    "# The builtin dialect is a collection of Operations, and Attributes that are expected\n",
    "# to be useful for most compilers, such as floating-point numbers, integers,\n",
    "# arrays, tensors, and more.\n",
    "from xdsl.dialects.builtin import DenseIntOrFPElementsAttr, Float64Type, TensorType, f64\n",
    "\n",
    "# The xdsl.ir module implements the things we've mentioned in this chapter,\n",
    "# especially the equivalents of MLIR concepts.\n",
    "# xdsl.irdl provides a declarative Python API for Operation definitions\n",
    "from xdsl.irdl import IRDLOperation, attr_def, irdl_op_definition, result_def\n",
    "from xdsl.utils.exceptions import VerifyException\n",
    "\n",
    "TensorTypeF64: TypeAlias = TensorType[Float64Type]\n",
    "\n",
    "\n",
    "# A decorator to help implement some methods required by xDSL\n",
    "@irdl_op_definition\n",
    "class ConstantOp(IRDLOperation):\n",
    "    \"\"\"\n",
    "    Constant operation turns a literal into an SSA value. The data is attached\n",
    "    to the operation as an attribute. For example:\n",
    "\n",
    "    ```mlir\n",
    "      %0 = \"toy.constant\"() {\"value\" = dense<[[1, 2, 3], [4, 5, 6]]> : tensor<2x3xf64>} : () -> tensor<2x3xf64>\n",
    "    ```\n",
    "    \"\"\"\n",
    "\n",
    "    # Every operation has a name. The format is `dialect_name`.`operation_name`\n",
    "    name = \"toy.constant\"\n",
    "\n",
    "    # Attributes are defined using OpAttr, and can specify a type constraint on the attribute\n",
    "    value = attr_def(DenseIntOrFPElementsAttr[Float64Type])\n",
    "\n",
    "    # The result type annotation uses `Annotated`, a type in the `typing` module that\n",
    "    # allows for runtime annotation of types with arbitrary values. xDSL leverages\n",
    "    # this annotation to populate a `verify()` method that will signal if there is\n",
    "    # a type mismatch during construction.\n",
    "    res = result_def(TensorTypeF64)\n",
    "\n",
    "    def __init__(self, value: DenseIntOrFPElementsAttr):\n",
    "        super().__init__(result_types=[value.type], attributes={\"value\": value})\n",
    "\n",
    "    # Operations can provide helper constructors for ease of use\n",
    "    @staticmethod\n",
    "    def from_list(data: list[float], shape: list[int]) -> ConstantOp:\n",
    "        value = DenseIntOrFPElementsAttr.from_list(TensorType(f64, shape), data)\n",
    "        return ConstantOp(value)\n",
    "\n",
    "    def verify_(self) -> None:\n",
    "        if not self.res.type == self.value.type:\n",
    "            raise VerifyException(\n",
    "                \"Expected value and result types to be equal: \"\n",
    "                f\"{self.res.type}, {self.value.type}\"\n",
    "            )\n",
    "\n",
    "\n",
    "Printer().print_op(ConstantOp.from_list([1, 2, 3, 4, 5, 6], [2, 3]))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "b37ce48c",
   "metadata": {},
   "source": [
    "\n",
    "### Constructor helper\n",
    "\n",
    "```\n",
    "    @staticmethod\n",
    "    def from_list(data: list[float], shape: list[int]):\n",
    "        value = DenseIntOrFPElementsAttr.tensor_from_list(data, f64, shape)\n",
    "\n",
    "        return ConstantOp.create(result_types=[value.type],\n",
    "                                 attributes={\"value\": value})\n",
    "```\n",
    "\n",
    "Operations tend to have helper methods for constructing them, that call into the generic\n",
    "constructors on the Operation class. In this case, the client passes in a flat python\n",
    "`list` of `float`s for the data, and a shape definition. These get converted to the\n",
    "attribute and result type that the `create` method expects as input.\n",
    "\n",
    "### Custom verifier\n",
    "\n",
    "``` python\n",
    "    def verify_(self) -> None:\n",
    "        resultType = self.res.type\n",
    "        value = self.value\n",
    "        if not isinstance(resultType, TensorType):\n",
    "            raise VerifyException(\"Expected result type to be `TensorTypeF64`\")\n",
    "\n",
    "        if not isinstance(value, DenseIntOrFPElementsAttr):\n",
    "            raise VerifyException(\n",
    "                \"Expected value type to be instance of `DenseIntOrFPElementsAttr`\"\n",
    "            )\n",
    "\n",
    "        if resultType.get_shape() != value.shape:\n",
    "            raise VerifyException(\n",
    "                \"Expected value and result to have the same shape\")\n",
    "```\n",
    "\n",
    "One thing to notice here is that all of our Toy operations are printed using the\n",
    "generic assembly format. This format is the one shown when breaking down\n",
    "`toy.transpose` at the beginning of this chapter. MLIR allows for operations to\n",
    "define their own custom assembly format, either or imperatively via C++. Defining a custom \n",
    "assembly format allows for tailoring the generated IR into something a bit more readable \n",
    "by removing a lot of the fluff that is required by the generic format. Let's walk through \n",
    "an example of an operation format that we would like to simplify.\n",
    "\n",
    "This capability will soon be added to xDSL also, and will be interoperable with the MLIR\n",
    "format definitions."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "c7078602",
   "metadata": {},
   "source": [
    "## Another Look at the Generated Toy IR\n",
    "\n",
    "Let's construct the same Toy program using the Python API:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "720752cf",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "builtin.module {\n",
      "  \"toy.func\"() ({\n",
      "    %0 = \"toy.constant\"() {value = dense<[[1.000000e+00, 2.000000e+00, 3.000000e+00], [4.000000e+00, 5.000000e+00, 6.000000e+00]]> : tensor<2x3xf64>} : () -> tensor<2x3xf64>\n",
      "    %1 = \"toy.reshape\"(%0) : (tensor<2x3xf64>) -> tensor<2x3xf64>\n",
      "    %2 = \"toy.constant\"() {value = dense<[1.000000e+00, 2.000000e+00, 3.000000e+00, 4.000000e+00, 5.000000e+00, 6.000000e+00]> : tensor<6xf64>} : () -> tensor<6xf64>\n",
      "    %3 = \"toy.reshape\"(%2) : (tensor<6xf64>) -> tensor<3x2xf64>\n",
      "    %4 = \"toy.reshape\"(%3) : (tensor<3x2xf64>) -> tensor<2x3xf64>\n",
      "    %5 = \"toy.add\"(%1, %4) : (tensor<2x3xf64>, tensor<2x3xf64>) -> tensor<2x3xf64>\n",
      "    \"toy.print\"(%5) : (tensor<2x3xf64>) -> ()\n",
      "    \"toy.return\"() : () -> ()\n",
      "  }) {sym_name = \"main\", function_type = () -> ()} : () -> ()\n",
      "}"
     ]
    }
   ],
   "source": [
    "from toy.dialects import toy\n",
    "\n",
    "from xdsl.builder import Builder\n",
    "from xdsl.dialects.builtin import FunctionType, ModuleOp\n",
    "\n",
    "\n",
    "@ModuleOp\n",
    "@Builder.implicit_region\n",
    "def module_op():\n",
    "    main_type = FunctionType.from_lists([], [])\n",
    "\n",
    "    @Builder.implicit_region\n",
    "    def main() -> None:\n",
    "        a_0 = toy.ConstantOp.from_list([1, 2, 3, 4, 5, 6], [2, 3]).res\n",
    "        a = toy.ReshapeOp(a_0, [2, 3]).res\n",
    "        b_0 = toy.ConstantOp.from_list([1, 2, 3, 4, 5, 6], [6]).res\n",
    "        b = toy.ReshapeOp(b_0, [3, 2]).res\n",
    "        c = toy.ReshapeOp(b, [2, 3]).res\n",
    "        d = toy.AddOp(a, c).res\n",
    "        toy.PrintOp(d)\n",
    "        toy.ReturnOp()\n",
    "\n",
    "    toy.FuncOp(\"main\", main_type, main)\n",
    "\n",
    "\n",
    "Printer().print_op(module_op)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "f2ce2423",
   "metadata": {},
   "source": [
    "First thing we see is that the whole program is wrapped in `builtin.module`. Modules\n",
    "roughly represent a parsed input file. Nested one level deep is a `toy.func`, representing\n",
    "the `main` function defined in the source. Inside it are a list of instructions.\n",
    "\n",
    "The next four lines of MLIR ops correspond to the first two lines of the Toy program.\n",
    "\n",
    "``` MLIR\n",
    "%0 = \"toy.constant\"() {\"value\" = dense<[[1, 2, 3], [4, 5, 6]]> : tensor<2x3xf64>} : () -> tensor<2x3xf64>\n",
    "%1 = \"toy.reshape\"(%0) : (tensor<2x3xf64>) -> tensor<2x3xf64>\n",
    "%2 = \"toy.constant\"() {\"value\" = dense<[1, 2, 3, 4, 5, 6]> : tensor<6xf64>} : () -> tensor<6xf64>\n",
    "%3 = \"toy.reshape\"(%2) : (tensor<6xf64>) -> tensor<3x2xf64>\n",
    "```\n",
    "\n",
    "``` Python\n",
    "var a<2, 3> = [[1, 2, 3], [4, 5, 6]];\n",
    "var b<3, 2> = [1, 2, 3, 4, 5, 6];\n",
    "```\n",
    "\n",
    "Because the types on the left of the `=` operator might not correspond to the type of the \n",
    "literal, the reshape operations are inserted. Most of the time the shapes will match,\n",
    "and the reshape will be redundant. Let's take a look at how to optimise our code to remove\n",
    "redundant reshapes using the xDSL infrastructure."
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
